#!/bin/bash

# 高级SSH远程命令执行脚本
# 配置文件: ~/.ssh/remote_exec.conf

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
PURPLE='\033[0;35m'
NC='\033[0m'

# 打印函数
print_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }
print_debug() { echo -e "${CYAN}[DEBUG]${NC} $1"; }
print_header() { echo -e "${PURPLE}[HEADER]${NC} $1"; }

CONFIG_FILE="$HOME/.ssh/remote_exec.conf"

# 默认配置
REMOTE_USER=""
REMOTE_HOST=""
REMOTE_PORT="22"
SSH_KEY=""
AUTO_CONFIRM="false"
SSH_OPTIONS="-o ConnectTimeout=10 -o ServerAliveInterval=60"

# 加载配置文件
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        print_info "加载配置文件: $CONFIG_FILE"
        source "$CONFIG_FILE"
    else
        create_default_config
    fi
}

# 创建默认配置文件
create_default_config() {
    print_warning "配置文件不存在，创建默认配置"
    cat > "$CONFIG_FILE" << EOF
# SSH远程执行配置
# 请修改以下配置为您的信息

REMOTE_USER="$(whoami)"
REMOTE_HOST="192.168.1.100"
REMOTE_PORT="22"
SSH_KEY="$HOME/.ssh/id_rsa"
AUTO_CONFIRM="false"
SSH_OPTIONS="-o ConnectTimeout=10 -o ServerAliveInterval=60"
EOF
    print_success "已创建默认配置文件: $CONFIG_FILE"
    print_info "请编辑该文件配置您的远程主机信息"
    exit 1
}

# 显示帮助信息
show_help() {
    print_header "SSH远程命令执行脚本"
    cat << EOF

用法: $0 [选项] [命令]

选项:
    -h, --help          显示此帮助信息
    -c, --config        编辑配置文件
    -i, --interactive   交互模式
    -f, --file [文件]   执行文件中的命令
    -y, --yes           自动确认，不提示
    -v, --verbose       详细输出模式
    --host [主机]       指定远程主机
    --user [用户]       指定用户名
    --port [端口]       指定SSH端口
    --key [密钥文件]    指定SSH密钥文件

示例:
    $0 "ls -l"
    $0 -i
    $0 -f commands.txt
    $0 --host 192.168.1.100 --user root "df -h"
    $0 -v "cat /etc/os-release"

配置文件位置: $CONFIG_FILE
EOF
}

# 检查必要工具
check_dependencies() {
    if ! command -v ssh &> /dev/null; then
        print_error "SSH客户端未安装，请先安装openssh-client"
        exit 1
    fi
}

# 测试SSH连接
test_ssh_connection() {
    print_info "测试SSH连接到 $REMOTE_USER@$REMOTE_HOST:$REMOTE_PORT..."
    
    local ssh_test_cmd="ssh $SSH_OPTIONS -p $REMOTE_PORT"
    if [ -n "$SSH_KEY" ] && [ -f "$SSH_KEY" ]; then
        ssh_test_cmd="$ssh_test_cmd -i $SSH_KEY"
    fi
    
    if $ssh_test_cmd -o ConnectTimeout=5 -q $REMOTE_USER@$REMOTE_HOST "echo '连接成功'"; then
        print_success "SSH连接测试成功"
        return 0
    else
        print_error "SSH连接测试失败"
        return 1
    fi
}

# 交互模式
interactive_mode() {
    print_header "SSH远程命令执行交互模式"
    
    # 测试连接
    if ! test_ssh_connection; then
        print_error "无法建立SSH连接，请检查配置"
        return 1
    fi
    
    # 获取远程主机信息
    if [ -z "$REMOTE_HOST" ]; then
        read -p "$(print_info "请输入远程主机IP或域名: ")" REMOTE_HOST
    fi
    
    if [ -z "$REMOTE_USER" ]; then
        read -p "$(print_info "请输入用户名: ")" REMOTE_USER
    fi
    
    print_success "已连接到 $REMOTE_USER@$REMOTE_HOST"
    print_info "输入 'quit', 'exit' 或按 Ctrl+C 退出"
    print_info "输入 'help' 查看可用命令"
    
    while true; do
        echo
        read -p "$(print_info "$REMOTE_USER@$REMOTE_HOST:~$ ")" CMD
        
        case "$CMD" in
            quit|exit)
                print_info "退出交互模式"
                break
                ;;
            help|?)
                cat << EOF
可用命令:
- quit, exit: 退出
- help, ?: 显示帮助
- hostname: 显示主机名
- uptime: 显示运行时间
- disk: 显示磁盘使用情况
- memory: 显示内存使用情况
- processes: 显示进程数
- custom: 执行任意命令
EOF
                ;;
            hostname)
                execute_command "hostname" "显示主机名"
                ;;
            uptime)
                execute_command "uptime" "显示系统运行时间"
                ;;
            disk)
                execute_command "df -h" "显示磁盘使用情况"
                ;;
            memory)
                execute_command "free -h" "显示内存使用情况"
                ;;
            processes)
                execute_command "ps aux | wc -l" "显示进程数量"
                ;;
            "")
                continue
                ;;
            *)
                execute_command "$CMD" "自定义命令"
                ;;
        esac
    done
}

# 执行命令
execute_command() {
    local cmd="$1"
    local description="$2"
    
    if [ -z "$cmd" ]; then
        print_error "命令不能为空"
        return 1
    fi
    
    if [ -n "$description" ]; then
        print_info "执行: $description"
    fi
    
    if [ "$AUTO_CONFIRM" != "true" ]; then
        print_warning "命令: $cmd"
        read -p "$(print_info "确认执行? (y/N): ")" confirm
        if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
            print_info "已取消"
            return
        fi
    fi
    
    echo -e "${BLUE}========================================${NC}"
    print_info "在 $REMOTE_USER@$REMOTE_HOST 上执行命令"
    print_debug "命令: $cmd"
    echo -e "${BLUE}========================================${NC}"
    
    # 构建SSH命令
    local ssh_cmd="ssh $SSH_OPTIONS"
    if [ -n "$SSH_KEY" ] && [ -f "$SSH_KEY" ]; then
        ssh_cmd="$ssh_cmd -i $SSH_KEY"
        print_debug "使用密钥: $SSH_KEY"
    fi
    if [ -n "$REMOTE_PORT" ]; then
        ssh_cmd="$ssh_cmd -p $REMOTE_PORT"
    fi
    
    # 执行命令
    $ssh_cmd $REMOTE_USER@$REMOTE_HOST "$cmd"
    
    local exit_code=$?
    echo -e "${BLUE}========================================${NC}"
    
    if [ $exit_code -eq 0 ]; then
        print_success "命令执行成功"
    else
        print_error "命令执行失败 (退出码: $exit_code)"
    fi
    
    return $exit_code
}

# 执行文件中的命令
execute_file() {
    local file="$1"
    
    if [ ! -f "$file" ]; then
        print_error "文件不存在: $file"
        exit 1
    fi
    
    print_info "从文件执行命令: $file"
    
    # 测试连接
    if ! test_ssh_connection; then
        print_error "无法建立SSH连接，请检查配置"
        exit 1
    fi
    
    local line_num=0
    local success_count=0
    local fail_count=0
    
    echo -e "${PURPLE}========================================${NC}"
    print_header "开始批量执行命令"
    echo -e "${PURPLE}========================================${NC}"
    
    while IFS= read -r line; do
        line_num=$((line_num + 1))
        
        # 跳过空行和注释
        if [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]]; then
            print_debug "跳过第 $line_num 行: $line"
            continue
        fi
        
        # 去除前后空格
        line=$(echo "$line" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
        
        print_info "执行第 $line_num 行: $line"
        execute_command "$line"
        
        if [ $? -eq 0 ]; then
            success_count=$((success_count + 1))
        else
            fail_count=$((fail_count + 1))
        fi
        
        echo
    done < "$file"
    
    echo -e "${PURPLE}========================================${NC}"
    print_header "批量执行完成"
    print_success "成功: $success_count 个命令"
    if [ $fail_count -gt 0 ]; then
        print_error "失败: $fail_count 个命令"
    fi
    print_info "总计: $line_num 行，执行了 $((success_count + fail_count)) 个命令"
    echo -e "${PURPLE}========================================${NC}"
}

# 显示连接信息
show_connection_info() {
    print_header "连接信息"
    print_info "远程主机: $REMOTE_USER@$REMOTE_HOST:$REMOTE_PORT"
    if [ -n "$SSH_KEY" ] && [ -f "$SSH_KEY" ]; then
        print_info "SSH密钥: $SSH_KEY"
    else
        print_info "认证方式: 密码认证"
    fi
    echo
}

# 主函数
main() {
    check_dependencies
    load_config
    
    local VERBOSE="false"
    
    # 解析参数
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)
                show_help
                exit 0
                ;;
            -c|--config)
                ${EDITOR:-nano} "$CONFIG_FILE"
                print_success "配置文件已编辑，请重新运行脚本"
                exit 0
                ;;
            -i|--interactive)
                show_connection_info
                interactive_mode
                exit 0
                ;;
            -f|--file)
                if [ -n "$2" ]; then
                    show_connection_info
                    execute_file "$2"
                    exit 0
                else
                    print_error "--file 参数需要指定文件路径"
                    exit 1
                fi
                ;;
            -y|--yes)
                AUTO_CONFIRM="true"
                shift
                ;;
            -v|--verbose)
                VERBOSE="true"
                SSH_OPTIONS="$SSH_OPTIONS -v"
                shift
                ;;
            --host)
                REMOTE_HOST="$2"
                shift 2
                ;;
            --user)
                REMOTE_USER="$2"
                shift 2
                ;;
            --port)
                REMOTE_PORT="$2"
                shift 2
                ;;
            --key)
                SSH_KEY="$2"
                shift 2
                ;;
            -*)
                print_error "未知选项: $1"
                show_help
                exit 1
                ;;
            *)
                CMD="$*"
                show_connection_info
                if ! test_ssh_connection; then
                    exit 1
                fi
                execute_command "$CMD"
                exit $?
                ;;
        esac
    done
    
    # 如果没有参数，显示帮助
    show_help
}

# 信号处理
trap 'print_error "脚本被用户中断"; exit 1' INT TERM

# 运行主函数
main "$@"
